{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Top1wCkMFu6C"
   },
   "source": [
    "# Fine-tune MiDashengLM with ESC-50\n",
    "\n",
    "## Pre-run checklist\n",
    "\n",
    "Before running, ensure MDL-Toolkit is properly installed. Running the command below should print the help message of `mdl-toolkit`. If it fails, please check your installation. For more details, see the Installation Guide at ./installation.md.\n",
    "\n",
    "> ### Note\n",
    "> When running locally, it is strongly recommended to install MDL-Toolkit into an isolated virtual environment to avoid dependency issues. In notebooks, environment handling can be tricky, so it’s fine to install MDL-Toolkit into the notebook’s current environment.\n",
    "\n",
    "> ### Note\n",
    "> During the example execution, the ESC-50 dataset and the full MiDashengLM-7B bf16 weights will be downloaded over the network. Ensure stable connections to GitHub and Hugging Face and enough disk space.\n",
    ">\n",
    "> You can also configure MDL-Toolkit to download models from Modelscope. To use Modelscope, make sure you installed the `modelscope` extra and add `--from-modelscope true` to `mdl-toolkit` commands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OWfIBxnmFwft",
    "outputId": "88d6c45b-c6aa-4dc5-c1c3-b45f8e7f80a4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "usage: mdl-toolkit [-h] {train,convert-dataset,inference} ...\r\n",
      "\r\n",
      "options:\r\n",
      "  -h, --help            show this help message and exit\r\n",
      "\r\n",
      "subcommands:\r\n",
      "  {train,convert-dataset,inference}\r\n",
      "    train\r\n",
      "    convert-dataset\r\n",
      "    inference\r\n"
     ]
    }
   ],
   "source": [
    "# Install MDL-Toolkit, for example:\n",
    "# !pip install mdl-toolkit\n",
    "!mdl-toolkit --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SWHSmQX9GLRl"
   },
   "source": [
    "## Data preparation\n",
    "\n",
    "### Download and extract the ESC-50 dataset\n",
    "\n",
    "Run the following command to download and extract the dataset. You can also obtain the dataset by other means; adjust paths in later steps accordingly.\n",
    "\n",
    "> ### Network access\n",
    "> This downloads the dataset (~615 MiB) from GitHub and may take some time. Please ensure your network is stable and you have sufficient storage space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KB3C95O3GMBz",
    "outputId": "66b519be-7055-4f9f-c566-52bd70513fd5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "File ‘ESC-50.zip’ already there; not retrieving.\r\n"
     ]
    }
   ],
   "source": [
    "!wget -nc https://github.com/karoldvl/ESC-50/archive/master.zip -O ESC-50.zip\n",
    "!unzip -o -q ESC-50.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KyN8OKDlHwkC"
   },
   "source": [
    "Now the ESC-50 dataset should be available under `ESC-50-master`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vunzt5yTH2ov"
   },
   "source": [
    "### Convert the dataset to the required training format\n",
    "\n",
    "ESC-50 provides a CSV list of samples with five folds `1` to `5`. The first 10 rows are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "q9NsOpJxH5cp",
    "outputId": "db4196c8-12c2-426a-ff46-14b72e9e92c6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "filename,fold,target,category,esc10,src_file,take\r\n",
      "1-100032-A-0.wav,1,0,dog,True,100032,A\r\n",
      "1-100038-A-14.wav,1,14,chirping_birds,False,100038,A\r\n",
      "1-100210-A-36.wav,1,36,vacuum_cleaner,False,100210,A\r\n",
      "1-100210-B-36.wav,1,36,vacuum_cleaner,False,100210,B\r\n",
      "1-101296-A-19.wav,1,19,thunderstorm,False,101296,A\r\n",
      "1-101296-B-19.wav,1,19,thunderstorm,False,101296,B\r\n",
      "1-101336-A-30.wav,1,30,door_wood_knock,False,101336,A\r\n",
      "1-101404-A-34.wav,1,34,can_opening,False,101404,A\r\n",
      "1-103298-A-9.wav,1,9,crow,False,103298,A\r\n",
      "1-103995-A-30.wav,1,30,door_wood_knock,False,103995,A\r\n"
     ]
    }
   ],
   "source": [
    "!head -n 11 ESC-50-master/meta/esc50.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_ywvoU4PIOsq"
   },
   "source": [
    "You can split the dataset into training and test by `fold` and format each row so the model predicts the category in the form `category: <category>, target: <target>`. Feel free to modify the code below, for example, make the model output JSON."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "YpYPE2jlHxPW"
   },
   "outputs": [],
   "source": [
    "import csv\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "esc50_base = Path(\"ESC-50-master\")\n",
    "meta_file = esc50_base / \"meta\" / \"esc50.csv\"\n",
    "train_output = Path(\"train.csv\")\n",
    "test_output = Path(\"test.csv\")\n",
    "\n",
    "with (\n",
    "    open(meta_file, \"r\") as meta,\n",
    "    open(train_output, \"w\") as train,\n",
    "    open(test_output, \"w\") as test,\n",
    "):\n",
    "    reader = csv.DictReader(meta)\n",
    "    train_writer = csv.DictWriter(train, fieldnames=[\"audio\", \"prediction\"])\n",
    "    test_writer = csv.DictWriter(test, fieldnames=[\"audio\", \"prediction\"])\n",
    "    train_writer.writeheader()\n",
    "    test_writer.writeheader()\n",
    "\n",
    "    for row in reader:\n",
    "        writer = train_writer if row[\"fold\"] != \"5\" else test_writer\n",
    "        writer.writerow(\n",
    "            {\n",
    "                \"audio\": os.fspath(esc50_base / \"audio\" / row[\"filename\"]),\n",
    "                \"prediction\": f\"category: {row['category']}, target: {row['target']}\",\n",
    "            }\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "aAuoiS-rIxkR",
    "outputId": "04a0e304-f7bd-4255-85d7-6e47ac847932"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==== Train split ====\n",
      "audio,prediction\n",
      "ESC-50-master/audio/1-100032-A-0.wav,\"category: dog, target: 0\"\n",
      "ESC-50-master/audio/1-100038-A-14.wav,\"category: chirping_birds, target: 14\"\n",
      "ESC-50-master/audio/1-100210-A-36.wav,\"category: vacuum_cleaner, target: 36\"\n",
      "ESC-50-master/audio/1-100210-B-36.wav,\"category: vacuum_cleaner, target: 36\"\n",
      "ESC-50-master/audio/1-101296-A-19.wav,\"category: thunderstorm, target: 19\"\n",
      "ESC-50-master/audio/1-101296-B-19.wav,\"category: thunderstorm, target: 19\"\n",
      "ESC-50-master/audio/1-101336-A-30.wav,\"category: door_wood_knock, target: 30\"\n",
      "ESC-50-master/audio/1-101404-A-34.wav,\"category: can_opening, target: 34\"\n",
      "ESC-50-master/audio/1-103298-A-9.wav,\"category: crow, target: 9\"\n",
      "ESC-50-master/audio/1-103995-A-30.wav,\"category: door_wood_knock, target: 30\"\n",
      "==== Test split  ====\n",
      "audio,prediction\n",
      "ESC-50-master/audio/5-103415-A-2.wav,\"category: pig, target: 2\"\n",
      "ESC-50-master/audio/5-103416-A-2.wav,\"category: pig, target: 2\"\n",
      "ESC-50-master/audio/5-103418-A-2.wav,\"category: pig, target: 2\"\n",
      "ESC-50-master/audio/5-103420-A-2.wav,\"category: pig, target: 2\"\n",
      "ESC-50-master/audio/5-103421-A-2.wav,\"category: pig, target: 2\"\n",
      "ESC-50-master/audio/5-103422-A-2.wav,\"category: pig, target: 2\"\n",
      "ESC-50-master/audio/5-117118-A-42.wav,\"category: siren, target: 42\"\n",
      "ESC-50-master/audio/5-117120-A-42.wav,\"category: siren, target: 42\"\n",
      "ESC-50-master/audio/5-117122-A-42.wav,\"category: siren, target: 42\"\n",
      "ESC-50-master/audio/5-117250-A-2.wav,\"category: pig, target: 2\"\n"
     ]
    }
   ],
   "source": [
    "!echo '==== Train split ===='\n",
    "!head -n 11 train.csv\n",
    "!echo '==== Test split  ===='\n",
    "!head -n 11 test.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6twsROtkeeb8"
   },
   "source": [
    "### Check the pretrained model’s output\n",
    "\n",
    "MiDashengLM doesn’t know the ESC-50 label space out of the box, so it may output incorrect categories. In this tutorial, we’ll fine-tune the model to align outputs with the expected format. In practice, carefully designed prompts or decoding constraints may improve outputs without fine-tuning—choose what fits your use case.\n",
    "\n",
    "MDL-Toolkit provides a convenient inference command to quickly run inference without writing code. The input format matches training, except the `prediction` column is optional; the command will write predictions into `prediction` and preserve other columns. Because the inference input format is compatible with training, we can reuse the `test.csv` generated above to observe the base model’s outputs.\n",
    "\n",
    "Arguments:\n",
    "- `--model-name mispeech/midashenglm-7b-bf16`: the Hugging Face repo or local path of the model to use.\n",
    "\n",
    "> ### Note\n",
    "> This tutorial uses bf16 weights to reduce download and disk usage. If you already have the full fp32 weights, you may use `--model-name mispeech/midashenglm-7b` instead of `--model-name mispeech/midashenglm-7b-bf16` to avoid re-downloading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ncZKkhfGk0Aj",
    "outputId": "2b8d04c1-994d-4ef6-b64b-9b571d040263"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section'}\n",
      "Loading checkpoint shards: 100%|██████████████████| 4/4 [00:04<00:00,  1.03s/it]\n",
      "Generating train split: 400 examples [00:00, 5910.05 examples/s]\n",
      "Processing dataset: 100%|██████████████| 400/400 [00:18<00:00, 22.22 examples/s]\n",
      "Batching examples (num_proc=32): 100%|█| 400/400 [00:03<00:00, 127.43 examples/s\n",
      "Inference: 100%|████████████████████████████████| 32/32 [00:32<00:00,  1.02s/it]\n",
      "audio,prediction\n",
      "ESC-50-master/audio/5-103415-A-2.wav,\"category: livestock, category_id: 1\"\n",
      "ESC-50-master/audio/5-103416-A-2.wav,\"category: music, category_id: 1\"\n",
      "ESC-50-master/audio/5-103418-A-2.wav,\"category: pig, category_id: 1\"\n",
      "ESC-50-master/audio/5-103420-A-2.wav,\"category: animal, category_id: 1\"\n",
      "ESC-50-master/audio/5-103421-A-2.wav,\"category: pig, category_id: 1\"\n",
      "ESC-50-master/audio/5-103422-A-2.wav,\"category: animal, category_id: 1\"\n",
      "ESC-50-master/audio/5-117118-A-42.wav,\"category: alarm, category_id: 1\"\n",
      "ESC-50-master/audio/5-117120-A-42.wav,\"category: alarm, category_id: 1\"\n",
      "ESC-50-master/audio/5-117122-A-42.wav,\"category: alarm, category_id: 1\"\n",
      "ESC-50-master/audio/5-117250-A-2.wav,\"category: animal, category_id: 1\"\n"
     ]
    }
   ],
   "source": [
    "!mdl-toolkit inference \\\n",
    "    test.csv \\\n",
    "    --system-prompt \"Output the predicted category in the format of category: <category>, category_id: <category_id>.\" \\\n",
    "    --output orig-output.csv \\\n",
    "    --model-name mispeech/midashenglm-7b-bf16\n",
    "! head -n 11 orig-output.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Uz_9MBs1JBd7"
   },
   "source": [
    "### Convert the data\n",
    "\n",
    "To speed up training, we convert the dataset ahead of time. The commands below convert the CSV files and set a simple system prompt. You can skip this step and convert during training; in that case, replace dataset paths in later steps with the CSV paths, and expect some conversion time before each run.\n",
    "\n",
    "> ### Network access\n",
    "> The following command will download the tokenizer from Hugging Face. Ensure your network is stable and wait patiently.\n",
    ">\n",
    "> To download from Modelscope, add `--from-modelscope true` to the commands and ensure you installed the `modelscope` extra.\n",
    "\n",
    "Arguments:\n",
    "- `train.csv`: path to the input CSV file.\n",
    "- `--output train-converted/`: output directory for the converted dataset; it will be created and existing content overwritten.\n",
    "- `--system-prompt ...`: a simple system prompt to guide the model’s behavior."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qvxAlVXkI4_D",
    "outputId": "e640487d-9edc-4d04-a959-4ca02a2a12a7"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating train split: 1600 examples [00:00, 27101.33 examples/s]\n",
      "Processing dataset: 100%|████████████| 1600/1600 [00:46<00:00, 34.10 examples/s]\n",
      "Deriving labels for training (num_proc=32): 100%|█| 1600/1600 [00:02<00:00, 667.\n",
      "Saving the dataset (2/2 shards): 100%|█| 1600/1600 [00:03<00:00, 522.67 examples\n",
      "Processing dataset: 100%|██████████████| 400/400 [00:16<00:00, 23.59 examples/s]\n",
      "Deriving labels for training (num_proc=32): 100%|█| 400/400 [00:02<00:00, 172.46\n",
      "Saving the dataset (1/1 shards): 100%|█| 400/400 [00:00<00:00, 1083.64 examples/\n"
     ]
    }
   ],
   "source": [
    "!mdl-toolkit convert-dataset \\\n",
    "    train.csv \\\n",
    "    --output train-converted/ \\\n",
    "    --system-prompt \"Output the predicted category in the format of category: <category>, category_id: <category_id>.\"\n",
    "!mdl-toolkit convert-dataset \\\n",
    "    test.csv \\\n",
    "    --output test-converted/ \\\n",
    "    --system-prompt \"Output the predicted category in the format of category: <category>, category_id: <category_id>.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FdazC9j9LRD9"
   },
   "source": [
    "## Train the model\n",
    "\n",
    "We want the model to classify audio while following a strict output format. Formatting is simple, so we use a modest LoRA rank and evaluate periodically. In this tutorial, we’ll use bf16 weights by default and evaluate more sparsely to save compute.\n",
    "\n",
    "> ### Network access\n",
    "> The command will download model weights from Hugging Face and may take some time. Ensure a stable network, enough storage space, and wait patiently.\n",
    ">\n",
    "> To download from Modelscope, add `--from-modelscope true` and ensure you installed the `modelscope` extra.\n",
    "\n",
    "> ### Note\n",
    "> We recommend training on a high-performance GPU for speed; MDL-Toolkit will automatically detect and use available GPUs. By default, training uses a single GPU. If you have multiple GPUs, see the Distributed Training Guide at ./distributed.md. Avoid CPU-only training as it will be very slow.\n",
    ">\n",
    "> To run with bf16 precision, you’ll need about 18 GiB of VRAM. If VRAM is limited, try adding `--quantization 8bit` or `--quantization 4bit` to quantize the model on load with bitsandbytes. Note that quantization may reduce capabilities and lead to suboptimal results.\n",
    "\n",
    "Arguments:\n",
    "- `--lora-rank 32`: set LoRA rank to 32. For more complex tasks, consider increasing the rank.\n",
    "- `--eval-steps 100`: evaluate every 100 steps.\n",
    "- `--train-dataset train-converted/`: training dataset path; you can also specify the CSV path to convert on the fly.\n",
    "- `--eval-dataset test-converted/`: evaluation dataset path; you can also specify the CSV path. If omitted, evaluation is skipped.\n",
    "- `--output output/`: output directory for checkpoints and results; it will be created and existing content overwritten."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "sHZXuZrwLR_l",
    "outputId": "48eeba04-844f-4375-8c72-38274d0de609"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Distributed: NO\n",
      "Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section'}\n",
      "Loading checkpoint shards: 100%|██████████████████| 4/4 [00:04<00:00,  1.01s/it]\n",
      "Model loaded with torch.bfloat16\n",
      "trainable params: 68,968,448 || all params: 8,350,708,352 || trainable%: 0.8259\n",
      "Peak VRAM during loading: 15.684 GiB\n",
      "  0%|                                                   | 0/200 [00:00<?, ?it/s]`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\n",
      "`loss_type=None` was set in the config but it is unrecognized. Using the default loss: `ForCausalLMLoss`.\n",
      "{'loss': 3.8798, 'grad_norm': 3.2964117527008057, 'learning_rate': 0.0001, 'epoch': 0.01}\n",
      "{'loss': 3.6867, 'grad_norm': 3.0418500900268555, 'learning_rate': 9.95e-05, 'epoch': 0.01}\n",
      "{'loss': 3.1348, 'grad_norm': 2.080457925796509, 'learning_rate': 9.900000000000001e-05, 'epoch': 0.01}\n",
      "{'loss': 2.8754, 'grad_norm': 2.0675864219665527, 'learning_rate': 9.850000000000001e-05, 'epoch': 0.02}\n",
      "{'loss': 2.5002, 'grad_norm': 2.108095169067383, 'learning_rate': 9.8e-05, 'epoch': 0.03}\n",
      "{'loss': 1.9693, 'grad_norm': 2.2269351482391357, 'learning_rate': 9.75e-05, 'epoch': 0.03}\n",
      "{'loss': 1.6182, 'grad_norm': 2.1482107639312744, 'learning_rate': 9.7e-05, 'epoch': 0.04}\n",
      "{'loss': 1.4915, 'grad_norm': 1.8752493858337402, 'learning_rate': 9.65e-05, 'epoch': 0.04}\n",
      "{'loss': 1.3654, 'grad_norm': 1.7094244956970215, 'learning_rate': 9.6e-05, 'epoch': 0.04}\n",
      "{'loss': 1.0916, 'grad_norm': 1.8022851943969727, 'learning_rate': 9.55e-05, 'epoch': 0.05}\n",
      "{'loss': 0.9587, 'grad_norm': 2.0512678623199463, 'learning_rate': 9.5e-05, 'epoch': 0.06}\n",
      "{'loss': 0.6848, 'grad_norm': 2.383748769760132, 'learning_rate': 9.449999999999999e-05, 'epoch': 0.06}\n",
      "{'loss': 0.6857, 'grad_norm': 1.3146827220916748, 'learning_rate': 9.4e-05, 'epoch': 0.07}\n",
      "{'loss': 0.4695, 'grad_norm': 1.1016651391983032, 'learning_rate': 9.350000000000001e-05, 'epoch': 0.07}\n",
      "{'loss': 0.4135, 'grad_norm': 1.3678333759307861, 'learning_rate': 9.300000000000001e-05, 'epoch': 0.07}\n",
      "{'loss': 0.4895, 'grad_norm': 1.2987140417099, 'learning_rate': 9.250000000000001e-05, 'epoch': 0.08}\n",
      "{'loss': 0.5315, 'grad_norm': 1.2354402542114258, 'learning_rate': 9.200000000000001e-05, 'epoch': 0.09}\n",
      "{'loss': 0.4365, 'grad_norm': 1.5113264322280884, 'learning_rate': 9.15e-05, 'epoch': 0.09}\n",
      "{'loss': 0.3974, 'grad_norm': 1.2639994621276855, 'learning_rate': 9.1e-05, 'epoch': 0.1}\n",
      "{'loss': 0.3503, 'grad_norm': 0.8667172789573669, 'learning_rate': 9.05e-05, 'epoch': 0.1}\n",
      "{'loss': 0.4957, 'grad_norm': 1.751583456993103, 'learning_rate': 9e-05, 'epoch': 0.1}\n",
      "{'loss': 0.3468, 'grad_norm': 1.2601685523986816, 'learning_rate': 8.950000000000001e-05, 'epoch': 0.11}\n",
      "{'loss': 0.4162, 'grad_norm': 1.393071174621582, 'learning_rate': 8.900000000000001e-05, 'epoch': 0.12}\n",
      "{'loss': 0.3701, 'grad_norm': 0.8844708204269409, 'learning_rate': 8.850000000000001e-05, 'epoch': 0.12}\n",
      "{'loss': 0.3662, 'grad_norm': 1.160038948059082, 'learning_rate': 8.800000000000001e-05, 'epoch': 0.12}\n",
      "{'loss': 0.3112, 'grad_norm': 1.4010485410690308, 'learning_rate': 8.75e-05, 'epoch': 0.13}\n",
      "{'loss': 0.3587, 'grad_norm': 1.1901799440383911, 'learning_rate': 8.7e-05, 'epoch': 0.14}\n",
      "{'loss': 0.2475, 'grad_norm': 0.8006367683410645, 'learning_rate': 8.65e-05, 'epoch': 0.14}\n",
      "{'loss': 0.3688, 'grad_norm': 1.2586493492126465, 'learning_rate': 8.6e-05, 'epoch': 0.14}\n",
      "{'loss': 0.2328, 'grad_norm': 0.8907356262207031, 'learning_rate': 8.55e-05, 'epoch': 0.15}\n",
      "{'loss': 0.3296, 'grad_norm': 0.9100168347358704, 'learning_rate': 8.5e-05, 'epoch': 0.15}\n",
      "{'loss': 0.2667, 'grad_norm': 1.0335599184036255, 'learning_rate': 8.450000000000001e-05, 'epoch': 0.16}\n",
      "{'loss': 0.2561, 'grad_norm': 1.4544076919555664, 'learning_rate': 8.4e-05, 'epoch': 0.17}\n",
      "{'loss': 0.373, 'grad_norm': 1.2631508111953735, 'learning_rate': 8.35e-05, 'epoch': 0.17}\n",
      "{'loss': 0.2704, 'grad_norm': 1.1197847127914429, 'learning_rate': 8.3e-05, 'epoch': 0.17}\n",
      "{'loss': 0.2308, 'grad_norm': 1.8249311447143555, 'learning_rate': 8.25e-05, 'epoch': 0.18}\n",
      "{'loss': 0.399, 'grad_norm': 1.444654941558838, 'learning_rate': 8.2e-05, 'epoch': 0.18}\n",
      "{'loss': 0.2524, 'grad_norm': 1.5981743335723877, 'learning_rate': 8.15e-05, 'epoch': 0.19}\n",
      "{'loss': 0.2327, 'grad_norm': 0.8596041202545166, 'learning_rate': 8.1e-05, 'epoch': 0.2}\n",
      "{'loss': 0.277, 'grad_norm': 1.198796033859253, 'learning_rate': 8.05e-05, 'epoch': 0.2}\n",
      "{'loss': 0.4313, 'grad_norm': 1.183451771736145, 'learning_rate': 8e-05, 'epoch': 0.2}\n",
      "{'loss': 0.366, 'grad_norm': 1.2554408311843872, 'learning_rate': 7.950000000000001e-05, 'epoch': 0.21}\n",
      "{'loss': 0.2234, 'grad_norm': 1.1534292697906494, 'learning_rate': 7.900000000000001e-05, 'epoch': 0.21}\n",
      "{'loss': 0.1708, 'grad_norm': 0.8630083203315735, 'learning_rate': 7.850000000000001e-05, 'epoch': 0.22}\n",
      "{'loss': 0.168, 'grad_norm': 1.134637713432312, 'learning_rate': 7.800000000000001e-05, 'epoch': 0.23}\n",
      "{'loss': 0.2355, 'grad_norm': 1.262500286102295, 'learning_rate': 7.75e-05, 'epoch': 0.23}\n",
      "{'loss': 0.2813, 'grad_norm': 2.729074716567993, 'learning_rate': 7.7e-05, 'epoch': 0.23}\n",
      "{'loss': 0.1438, 'grad_norm': 1.3244102001190186, 'learning_rate': 7.65e-05, 'epoch': 0.24}\n",
      "{'loss': 0.1912, 'grad_norm': 1.1869083642959595, 'learning_rate': 7.6e-05, 'epoch': 0.24}\n",
      "{'loss': 0.2249, 'grad_norm': 1.447842001914978, 'learning_rate': 7.55e-05, 'epoch': 0.25}\n",
      "{'loss': 0.2473, 'grad_norm': 2.019524574279785, 'learning_rate': 7.500000000000001e-05, 'epoch': 0.26}\n",
      "{'loss': 0.1433, 'grad_norm': 1.6164042949676514, 'learning_rate': 7.450000000000001e-05, 'epoch': 0.26}\n",
      "{'loss': 0.1629, 'grad_norm': 0.9746866226196289, 'learning_rate': 7.4e-05, 'epoch': 0.27}\n",
      "{'loss': 0.1818, 'grad_norm': 1.499971866607666, 'learning_rate': 7.35e-05, 'epoch': 0.27}\n",
      "{'loss': 0.1286, 'grad_norm': 2.2042994499206543, 'learning_rate': 7.3e-05, 'epoch': 0.28}\n",
      "{'loss': 0.187, 'grad_norm': 1.9872567653656006, 'learning_rate': 7.25e-05, 'epoch': 0.28}\n",
      "{'loss': 0.1925, 'grad_norm': 0.9911997318267822, 'learning_rate': 7.2e-05, 'epoch': 0.28}\n",
      "{'loss': 0.2073, 'grad_norm': 1.5939922332763672, 'learning_rate': 7.15e-05, 'epoch': 0.29}\n",
      "{'loss': 0.0969, 'grad_norm': 0.7479377388954163, 'learning_rate': 7.1e-05, 'epoch': 0.29}\n",
      "{'loss': 0.1381, 'grad_norm': 0.6693164110183716, 'learning_rate': 7.05e-05, 'epoch': 0.3}\n",
      "{'loss': 0.1111, 'grad_norm': 1.0318104028701782, 'learning_rate': 7e-05, 'epoch': 0.3}\n",
      "{'loss': 0.1651, 'grad_norm': 1.8631339073181152, 'learning_rate': 6.95e-05, 'epoch': 0.31}\n",
      "{'loss': 0.115, 'grad_norm': 1.9326995611190796, 'learning_rate': 6.9e-05, 'epoch': 0.32}\n",
      "{'loss': 0.0707, 'grad_norm': 0.8733207583427429, 'learning_rate': 6.850000000000001e-05, 'epoch': 0.32}\n",
      "{'loss': 0.2227, 'grad_norm': 2.1967313289642334, 'learning_rate': 6.800000000000001e-05, 'epoch': 0.33}\n",
      "{'loss': 0.1862, 'grad_norm': 1.6454542875289917, 'learning_rate': 6.750000000000001e-05, 'epoch': 0.33}\n",
      "{'loss': 0.0961, 'grad_norm': 2.025951385498047, 'learning_rate': 6.7e-05, 'epoch': 0.34}\n",
      "{'loss': 0.1246, 'grad_norm': 1.2836817502975464, 'learning_rate': 6.65e-05, 'epoch': 0.34}\n",
      "{'loss': 0.0544, 'grad_norm': 0.6507941484451294, 'learning_rate': 6.6e-05, 'epoch': 0.34}\n",
      "{'loss': 0.0783, 'grad_norm': 1.6159645318984985, 'learning_rate': 6.55e-05, 'epoch': 0.35}\n",
      "{'loss': 0.0336, 'grad_norm': 0.5483852624893188, 'learning_rate': 6.500000000000001e-05, 'epoch': 0.35}\n",
      "{'loss': 0.0889, 'grad_norm': 1.3239760398864746, 'learning_rate': 6.450000000000001e-05, 'epoch': 0.36}\n",
      "{'loss': 0.1004, 'grad_norm': 1.406546950340271, 'learning_rate': 6.400000000000001e-05, 'epoch': 0.36}\n",
      "{'loss': 0.0557, 'grad_norm': 0.8986234664916992, 'learning_rate': 6.35e-05, 'epoch': 0.37}\n",
      "{'loss': 0.0254, 'grad_norm': 0.6604735851287842, 'learning_rate': 6.3e-05, 'epoch': 0.38}\n",
      "{'loss': 0.1176, 'grad_norm': 1.7354787588119507, 'learning_rate': 6.25e-05, 'epoch': 0.38}\n",
      "{'loss': 0.0666, 'grad_norm': 0.8596743941307068, 'learning_rate': 6.2e-05, 'epoch': 0.39}\n",
      "{'loss': 0.0532, 'grad_norm': 0.7732245326042175, 'learning_rate': 6.15e-05, 'epoch': 0.39}\n",
      "{'loss': 0.068, 'grad_norm': 1.0519976615905762, 'learning_rate': 6.1e-05, 'epoch': 0.4}\n",
      "{'loss': 0.0482, 'grad_norm': 0.8593562245368958, 'learning_rate': 6.05e-05, 'epoch': 0.4}\n",
      "{'loss': 0.0998, 'grad_norm': 2.5991556644439697, 'learning_rate': 6e-05, 'epoch': 0.41}\n",
      "{'loss': 0.0312, 'grad_norm': 1.1342809200286865, 'learning_rate': 5.95e-05, 'epoch': 0.41}\n",
      "{'loss': 0.045, 'grad_norm': 1.2447174787521362, 'learning_rate': 5.9e-05, 'epoch': 0.41}\n",
      "{'loss': 0.0338, 'grad_norm': 0.483461856842041, 'learning_rate': 5.85e-05, 'epoch': 0.42}\n",
      "{'loss': 0.1878, 'grad_norm': 1.7174185514450073, 'learning_rate': 5.8e-05, 'epoch': 0.42}\n",
      "{'loss': 0.0596, 'grad_norm': 2.2838027477264404, 'learning_rate': 5.7499999999999995e-05, 'epoch': 0.43}\n",
      "{'loss': 0.0406, 'grad_norm': 0.5791202187538147, 'learning_rate': 5.6999999999999996e-05, 'epoch': 0.43}\n",
      "{'loss': 0.0168, 'grad_norm': 0.40813687443733215, 'learning_rate': 5.65e-05, 'epoch': 0.44}\n",
      "{'loss': 0.0079, 'grad_norm': 0.5135295987129211, 'learning_rate': 5.6000000000000006e-05, 'epoch': 0.45}\n",
      "{'loss': 0.0256, 'grad_norm': 0.38277697563171387, 'learning_rate': 5.550000000000001e-05, 'epoch': 0.45}\n",
      "{'loss': 0.0787, 'grad_norm': 1.1708821058273315, 'learning_rate': 5.500000000000001e-05, 'epoch': 0.46}\n",
      "{'loss': 0.0056, 'grad_norm': 0.6055694222450256, 'learning_rate': 5.45e-05, 'epoch': 0.46}\n",
      "{'loss': 0.0353, 'grad_norm': 0.6405754685401917, 'learning_rate': 5.4000000000000005e-05, 'epoch': 0.47}\n",
      "{'loss': 0.0458, 'grad_norm': 0.5267824530601501, 'learning_rate': 5.3500000000000006e-05, 'epoch': 0.47}\n",
      "{'loss': 0.1109, 'grad_norm': 0.9342306852340698, 'learning_rate': 5.300000000000001e-05, 'epoch': 0.47}\n",
      "{'loss': 0.0253, 'grad_norm': 2.990312099456787, 'learning_rate': 5.25e-05, 'epoch': 0.48}\n",
      "{'loss': 0.0477, 'grad_norm': 0.6595431566238403, 'learning_rate': 5.2000000000000004e-05, 'epoch': 0.48}\n",
      "{'loss': 0.0278, 'grad_norm': 2.436335802078247, 'learning_rate': 5.1500000000000005e-05, 'epoch': 0.49}\n",
      "{'loss': 0.1186, 'grad_norm': 1.9407051801681519, 'learning_rate': 5.1000000000000006e-05, 'epoch': 0.49}\n",
      "{'loss': 0.0199, 'grad_norm': 0.434623658657074, 'learning_rate': 5.05e-05, 'epoch': 0.5}\n",
      " 50%|████████████████████▌                    | 100/200 [00:54<00:56,  1.77it/s]\n",
      "  0%|                                                    | 0/50 [00:00<?, ?it/s]\u001b[A\n",
      "  4%|█▊                                          | 2/50 [00:00<00:06,  7.25it/s]\u001b[A\n",
      "  6%|██▋                                         | 3/50 [00:00<00:09,  5.10it/s]\u001b[A\n",
      "  8%|███▌                                        | 4/50 [00:00<00:10,  4.31it/s]\u001b[A\n",
      " 10%|████▍                                       | 5/50 [00:01<00:11,  4.04it/s]\u001b[A\n",
      " 12%|█████▎                                      | 6/50 [00:01<00:11,  3.89it/s]\u001b[A\n",
      " 14%|██████▏                                     | 7/50 [00:01<00:11,  3.79it/s]\u001b[A\n",
      " 16%|███████                                     | 8/50 [00:01<00:11,  3.64it/s]\u001b[A\n",
      " 18%|███████▉                                    | 9/50 [00:02<00:11,  3.61it/s]\u001b[A\n",
      " 20%|████████▌                                  | 10/50 [00:02<00:11,  3.61it/s]\u001b[A\n",
      " 22%|█████████▍                                 | 11/50 [00:02<00:10,  3.59it/s]\u001b[A\n",
      " 24%|██████████▎                                | 12/50 [00:03<00:10,  3.61it/s]\u001b[A\n",
      " 26%|███████████▏                               | 13/50 [00:03<00:10,  3.62it/s]\u001b[A\n",
      " 28%|████████████                               | 14/50 [00:03<00:09,  3.63it/s]\u001b[A\n",
      " 30%|████████████▉                              | 15/50 [00:03<00:09,  3.55it/s]\u001b[A\n",
      " 32%|█████████████▊                             | 16/50 [00:04<00:09,  3.58it/s]\u001b[A\n",
      " 34%|██████████████▌                            | 17/50 [00:04<00:09,  3.60it/s]\u001b[A\n",
      " 36%|███████████████▍                           | 18/50 [00:04<00:09,  3.49it/s]\u001b[A\n",
      " 38%|████████████████▎                          | 19/50 [00:05<00:08,  3.49it/s]\u001b[A\n",
      " 40%|█████████████████▏                         | 20/50 [00:05<00:08,  3.53it/s]\u001b[A\n",
      " 42%|██████████████████                         | 21/50 [00:05<00:08,  3.53it/s]\u001b[A\n",
      " 44%|██████████████████▉                        | 22/50 [00:05<00:07,  3.52it/s]\u001b[A\n",
      " 46%|███████████████████▊                       | 23/50 [00:06<00:07,  3.55it/s]\u001b[A\n",
      " 48%|████████████████████▋                      | 24/50 [00:06<00:07,  3.54it/s]\u001b[A\n",
      " 50%|█████████████████████▌                     | 25/50 [00:06<00:07,  3.54it/s]\u001b[A\n",
      " 52%|██████████████████████▎                    | 26/50 [00:07<00:06,  3.49it/s]\u001b[A\n",
      " 54%|███████████████████████▏                   | 27/50 [00:07<00:06,  3.54it/s]\u001b[A\n",
      " 56%|████████████████████████                   | 28/50 [00:07<00:06,  3.56it/s]\u001b[A\n",
      " 58%|████████████████████████▉                  | 29/50 [00:07<00:05,  3.51it/s]\u001b[A\n",
      " 60%|█████████████████████████▊                 | 30/50 [00:08<00:05,  3.53it/s]\u001b[A\n",
      " 62%|██████████████████████████▋                | 31/50 [00:08<00:05,  3.56it/s]\u001b[A\n",
      " 64%|███████████████████████████▌               | 32/50 [00:08<00:05,  3.52it/s]\u001b[A\n",
      " 66%|████████████████████████████▍              | 33/50 [00:09<00:04,  3.53it/s]\u001b[A\n",
      " 68%|█████████████████████████████▏             | 34/50 [00:09<00:04,  3.53it/s]\u001b[A\n",
      " 70%|██████████████████████████████             | 35/50 [00:09<00:04,  3.54it/s]\u001b[A\n",
      " 72%|██████████████████████████████▉            | 36/50 [00:09<00:03,  3.50it/s]\u001b[A\n",
      " 74%|███████████████████████████████▊           | 37/50 [00:10<00:03,  3.46it/s]\u001b[A\n",
      " 76%|████████████████████████████████▋          | 38/50 [00:10<00:03,  3.49it/s]\u001b[A\n",
      " 78%|█████████████████████████████████▌         | 39/50 [00:10<00:03,  3.49it/s]\u001b[A\n",
      " 80%|██████████████████████████████████▍        | 40/50 [00:11<00:02,  3.49it/s]\u001b[A\n",
      " 82%|███████████████████████████████████▎       | 41/50 [00:11<00:02,  3.52it/s]\u001b[A\n",
      " 84%|████████████████████████████████████       | 42/50 [00:11<00:02,  3.55it/s]\u001b[A\n",
      " 86%|████████████████████████████████████▉      | 43/50 [00:11<00:01,  3.52it/s]\u001b[A\n",
      " 88%|█████████████████████████████████████▊     | 44/50 [00:12<00:01,  3.54it/s]\u001b[A\n",
      " 90%|██████████████████████████████████████▋    | 45/50 [00:12<00:01,  3.52it/s]\u001b[A\n",
      " 92%|███████████████████████████████████████▌   | 46/50 [00:12<00:01,  3.52it/s]\u001b[A\n",
      " 94%|████████████████████████████████████████▍  | 47/50 [00:13<00:00,  3.51it/s]\u001b[A\n",
      " 96%|█████████████████████████████████████████▎ | 48/50 [00:13<00:00,  3.56it/s]\u001b[A\n",
      "                                                                                \n",
      "\u001b[A{'eval_loss': 0.03699319064617157, 'eval_runtime': 14.1762, 'eval_samples_per_second': 28.216, 'eval_steps_per_second': 3.527, 'epoch': 0.5}\n",
      " 50%|████████████████████▌                    | 100/200 [01:09<00:56,  1.77it/s]\n",
      "100%|███████████████████████████████████████████| 50/50 [00:13<00:00,  3.59it/s]\u001b[A\n",
      "{'loss': 0.0089, 'grad_norm': 0.23225542902946472, 'learning_rate': 5e-05, 'epoch': 0.51}\n",
      "{'loss': 0.0913, 'grad_norm': 1.179492473602295, 'learning_rate': 4.9500000000000004e-05, 'epoch': 0.51}\n",
      "{'loss': 0.0633, 'grad_norm': 1.0864746570587158, 'learning_rate': 4.9e-05, 'epoch': 0.52}\n",
      "{'loss': 0.0453, 'grad_norm': 0.6859449148178101, 'learning_rate': 4.85e-05, 'epoch': 0.52}\n",
      "{'loss': 0.016, 'grad_norm': 0.3083536922931671, 'learning_rate': 4.8e-05, 'epoch': 0.53}\n",
      "{'loss': 0.0033, 'grad_norm': 0.10332754254341125, 'learning_rate': 4.75e-05, 'epoch': 0.53}\n",
      "{'loss': 0.0053, 'grad_norm': 0.19334560632705688, 'learning_rate': 4.7e-05, 'epoch': 0.54}\n",
      "{'loss': 0.0247, 'grad_norm': 0.4422135353088379, 'learning_rate': 4.6500000000000005e-05, 'epoch': 0.54}\n",
      "{'loss': 0.0234, 'grad_norm': 0.4461279511451721, 'learning_rate': 4.600000000000001e-05, 'epoch': 0.55}\n",
      "{'loss': 0.013, 'grad_norm': 0.4381876289844513, 'learning_rate': 4.55e-05, 'epoch': 0.55}\n",
      "{'loss': 0.0296, 'grad_norm': 0.39517849683761597, 'learning_rate': 4.5e-05, 'epoch': 0.56}\n",
      "{'loss': 0.004, 'grad_norm': 0.10346654057502747, 'learning_rate': 4.4500000000000004e-05, 'epoch': 0.56}\n",
      "{'loss': 0.0116, 'grad_norm': 0.30043676495552063, 'learning_rate': 4.4000000000000006e-05, 'epoch': 0.56}\n",
      "{'loss': 0.0116, 'grad_norm': 0.2605663537979126, 'learning_rate': 4.35e-05, 'epoch': 0.57}\n",
      "{'loss': 0.0064, 'grad_norm': 0.3289777934551239, 'learning_rate': 4.3e-05, 'epoch': 0.57}\n",
      "{'loss': 0.0409, 'grad_norm': 0.6381182074546814, 'learning_rate': 4.25e-05, 'epoch': 0.58}\n",
      "{'loss': 0.0094, 'grad_norm': 0.3302237391471863, 'learning_rate': 4.2e-05, 'epoch': 0.58}\n",
      "{'loss': 0.0185, 'grad_norm': 0.35747629404067993, 'learning_rate': 4.15e-05, 'epoch': 0.59}\n",
      "{'loss': 0.0164, 'grad_norm': 0.47557777166366577, 'learning_rate': 4.1e-05, 'epoch': 0.59}\n",
      "{'loss': 0.0135, 'grad_norm': 1.1637630462646484, 'learning_rate': 4.05e-05, 'epoch': 0.6}\n",
      "{'loss': 0.0278, 'grad_norm': 0.6992457509040833, 'learning_rate': 4e-05, 'epoch': 0.6}\n",
      "{'loss': 0.0014, 'grad_norm': 0.0339697040617466, 'learning_rate': 3.9500000000000005e-05, 'epoch': 0.61}\n",
      "{'loss': 0.0144, 'grad_norm': 0.40721485018730164, 'learning_rate': 3.9000000000000006e-05, 'epoch': 0.61}\n",
      "{'loss': 0.0021, 'grad_norm': 0.06161792203783989, 'learning_rate': 3.85e-05, 'epoch': 0.62}\n",
      "{'loss': 0.0071, 'grad_norm': 0.16749635338783264, 'learning_rate': 3.8e-05, 'epoch': 0.62}\n",
      "{'loss': 0.015, 'grad_norm': 0.3915780186653137, 'learning_rate': 3.7500000000000003e-05, 'epoch': 0.63}\n",
      "{'loss': 0.0065, 'grad_norm': 0.22043326497077942, 'learning_rate': 3.7e-05, 'epoch': 0.64}\n",
      "{'loss': 0.0016, 'grad_norm': 0.060331325978040695, 'learning_rate': 3.65e-05, 'epoch': 0.64}\n",
      "{'loss': 0.0313, 'grad_norm': 0.4154261350631714, 'learning_rate': 3.6e-05, 'epoch': 0.65}\n",
      "{'loss': 0.0043, 'grad_norm': 0.18240374326705933, 'learning_rate': 3.55e-05, 'epoch': 0.65}\n",
      "{'loss': 0.0462, 'grad_norm': 0.5258132815361023, 'learning_rate': 3.5e-05, 'epoch': 0.66}\n",
      "{'loss': 0.0034, 'grad_norm': 0.12947408854961395, 'learning_rate': 3.45e-05, 'epoch': 0.66}\n",
      "{'loss': 0.0005, 'grad_norm': 0.01885071024298668, 'learning_rate': 3.4000000000000007e-05, 'epoch': 0.67}\n",
      "{'loss': 0.0133, 'grad_norm': 0.6813490986824036, 'learning_rate': 3.35e-05, 'epoch': 0.67}\n",
      "{'loss': 0.0011, 'grad_norm': 0.023346586152911186, 'learning_rate': 3.3e-05, 'epoch': 0.68}\n",
      "{'loss': 0.0578, 'grad_norm': 0.8785558938980103, 'learning_rate': 3.2500000000000004e-05, 'epoch': 0.68}\n",
      "{'loss': 0.001, 'grad_norm': 0.03433295339345932, 'learning_rate': 3.2000000000000005e-05, 'epoch': 0.69}\n",
      "{'loss': 0.0022, 'grad_norm': 0.07163651287555695, 'learning_rate': 3.15e-05, 'epoch': 0.69}\n",
      "{'loss': 0.002, 'grad_norm': 0.05869655683636665, 'learning_rate': 3.1e-05, 'epoch': 0.69}\n",
      "{'loss': 0.0033, 'grad_norm': 0.1852327436208725, 'learning_rate': 3.05e-05, 'epoch': 0.7}\n",
      "{'loss': 0.0084, 'grad_norm': 0.3363635540008545, 'learning_rate': 3e-05, 'epoch': 0.7}\n",
      "{'loss': 0.0212, 'grad_norm': 0.5239258408546448, 'learning_rate': 2.95e-05, 'epoch': 0.71}\n",
      "{'loss': 0.0156, 'grad_norm': 0.509312629699707, 'learning_rate': 2.9e-05, 'epoch': 0.71}\n",
      "{'loss': 0.0013, 'grad_norm': 0.05449633672833443, 'learning_rate': 2.8499999999999998e-05, 'epoch': 0.72}\n",
      "{'loss': 0.0361, 'grad_norm': 0.5985652804374695, 'learning_rate': 2.8000000000000003e-05, 'epoch': 0.72}\n",
      "{'loss': 0.0003, 'grad_norm': 0.010093459859490395, 'learning_rate': 2.7500000000000004e-05, 'epoch': 0.73}\n",
      "{'loss': 0.0075, 'grad_norm': 0.3111060559749603, 'learning_rate': 2.7000000000000002e-05, 'epoch': 0.73}\n",
      "{'loss': 0.0161, 'grad_norm': 0.40252572298049927, 'learning_rate': 2.6500000000000004e-05, 'epoch': 0.74}\n",
      "{'loss': 0.0044, 'grad_norm': 0.2180543839931488, 'learning_rate': 2.6000000000000002e-05, 'epoch': 0.74}\n",
      "{'loss': 0.0009, 'grad_norm': 0.02918129600584507, 'learning_rate': 2.5500000000000003e-05, 'epoch': 0.75}\n",
      "{'loss': 0.0013, 'grad_norm': 0.058701030910015106, 'learning_rate': 2.5e-05, 'epoch': 0.76}\n",
      "{'loss': 0.0015, 'grad_norm': 0.06371690332889557, 'learning_rate': 2.45e-05, 'epoch': 0.76}\n",
      "{'loss': 0.0289, 'grad_norm': 0.7292153239250183, 'learning_rate': 2.4e-05, 'epoch': 0.77}\n",
      "{'loss': 0.0028, 'grad_norm': 0.09709018468856812, 'learning_rate': 2.35e-05, 'epoch': 0.77}\n",
      "{'loss': 0.0173, 'grad_norm': 0.49819281697273254, 'learning_rate': 2.3000000000000003e-05, 'epoch': 0.78}\n",
      "{'loss': 0.0021, 'grad_norm': 0.08634943515062332, 'learning_rate': 2.25e-05, 'epoch': 0.78}\n",
      "{'loss': 0.0026, 'grad_norm': 0.1813768446445465, 'learning_rate': 2.2000000000000003e-05, 'epoch': 0.79}\n",
      "{'loss': 0.0024, 'grad_norm': 0.1100187748670578, 'learning_rate': 2.15e-05, 'epoch': 0.79}\n",
      "{'loss': 0.0015, 'grad_norm': 0.041948121041059494, 'learning_rate': 2.1e-05, 'epoch': 0.8}\n",
      "{'loss': 0.0003, 'grad_norm': 0.010011507198214531, 'learning_rate': 2.05e-05, 'epoch': 0.8}\n",
      "{'loss': 0.0087, 'grad_norm': 0.30718350410461426, 'learning_rate': 2e-05, 'epoch': 0.81}\n",
      "{'loss': 0.1002, 'grad_norm': 0.9224221706390381, 'learning_rate': 1.9500000000000003e-05, 'epoch': 0.81}\n",
      "{'loss': 0.0005, 'grad_norm': 0.011980859562754631, 'learning_rate': 1.9e-05, 'epoch': 0.81}\n",
      "{'loss': 0.0244, 'grad_norm': 0.6678000688552856, 'learning_rate': 1.85e-05, 'epoch': 0.82}\n",
      "{'loss': 0.0206, 'grad_norm': 0.49770551919937134, 'learning_rate': 1.8e-05, 'epoch': 0.82}\n",
      "{'loss': 0.0013, 'grad_norm': 0.08127743750810623, 'learning_rate': 1.75e-05, 'epoch': 0.83}\n",
      "{'loss': 0.0164, 'grad_norm': 0.5039976835250854, 'learning_rate': 1.7000000000000003e-05, 'epoch': 0.83}\n",
      "{'loss': 0.0038, 'grad_norm': 0.15809063613414764, 'learning_rate': 1.65e-05, 'epoch': 0.84}\n",
      "{'loss': 0.0251, 'grad_norm': 0.506536602973938, 'learning_rate': 1.6000000000000003e-05, 'epoch': 0.84}\n",
      "{'loss': 0.0017, 'grad_norm': 0.05380581319332123, 'learning_rate': 1.55e-05, 'epoch': 0.85}\n",
      "{'loss': 0.0656, 'grad_norm': 0.9112598299980164, 'learning_rate': 1.5e-05, 'epoch': 0.85}\n",
      "{'loss': 0.0014, 'grad_norm': 0.05966261029243469, 'learning_rate': 1.45e-05, 'epoch': 0.86}\n",
      "{'loss': 0.0004, 'grad_norm': 0.03510352596640587, 'learning_rate': 1.4000000000000001e-05, 'epoch': 0.86}\n",
      "{'loss': 0.0161, 'grad_norm': 0.6491398215293884, 'learning_rate': 1.3500000000000001e-05, 'epoch': 0.87}\n",
      "{'loss': 0.0006, 'grad_norm': 0.016234172508120537, 'learning_rate': 1.3000000000000001e-05, 'epoch': 0.88}\n",
      "{'loss': 0.0044, 'grad_norm': 0.1581207662820816, 'learning_rate': 1.25e-05, 'epoch': 0.88}\n",
      "{'loss': 0.0008, 'grad_norm': 0.04637967795133591, 'learning_rate': 1.2e-05, 'epoch': 0.89}\n",
      "{'loss': 0.0008, 'grad_norm': 0.01929531618952751, 'learning_rate': 1.1500000000000002e-05, 'epoch': 0.89}\n",
      "{'loss': 0.0009, 'grad_norm': 0.03805166110396385, 'learning_rate': 1.1000000000000001e-05, 'epoch': 0.9}\n",
      "{'loss': 0.0015, 'grad_norm': 0.04581131786108017, 'learning_rate': 1.05e-05, 'epoch': 0.9}\n",
      "{'loss': 0.0031, 'grad_norm': 0.0859525129199028, 'learning_rate': 1e-05, 'epoch': 0.91}\n",
      "{'loss': 0.0064, 'grad_norm': 0.3097107708454132, 'learning_rate': 9.5e-06, 'epoch': 0.91}\n",
      "{'loss': 0.0135, 'grad_norm': 0.3599609136581421, 'learning_rate': 9e-06, 'epoch': 0.92}\n",
      "{'loss': 0.0188, 'grad_norm': 0.5517334938049316, 'learning_rate': 8.500000000000002e-06, 'epoch': 0.92}\n",
      "{'loss': 0.0061, 'grad_norm': 0.32371991872787476, 'learning_rate': 8.000000000000001e-06, 'epoch': 0.93}\n",
      "{'loss': 0.0146, 'grad_norm': 0.5317452549934387, 'learning_rate': 7.5e-06, 'epoch': 0.93}\n",
      "{'loss': 0.0007, 'grad_norm': 0.017851119861006737, 'learning_rate': 7.000000000000001e-06, 'epoch': 0.94}\n",
      "{'loss': 0.0032, 'grad_norm': 0.12277483940124512, 'learning_rate': 6.5000000000000004e-06, 'epoch': 0.94}\n",
      "{'loss': 0.001, 'grad_norm': 0.038719117641448975, 'learning_rate': 6e-06, 'epoch': 0.94}\n",
      "{'loss': 0.0464, 'grad_norm': 1.0585988759994507, 'learning_rate': 5.500000000000001e-06, 'epoch': 0.95}\n",
      "{'loss': 0.0038, 'grad_norm': 0.11854031682014465, 'learning_rate': 5e-06, 'epoch': 0.95}\n",
      "{'loss': 0.0134, 'grad_norm': 0.44210073351860046, 'learning_rate': 4.5e-06, 'epoch': 0.96}\n",
      "{'loss': 0.0228, 'grad_norm': 0.669529914855957, 'learning_rate': 4.000000000000001e-06, 'epoch': 0.96}\n",
      "{'loss': 0.0016, 'grad_norm': 0.06973986327648163, 'learning_rate': 3.5000000000000004e-06, 'epoch': 0.97}\n",
      "{'loss': 0.0007, 'grad_norm': 0.018438400700688362, 'learning_rate': 3e-06, 'epoch': 0.97}\n",
      "{'loss': 0.0033, 'grad_norm': 0.14516839385032654, 'learning_rate': 2.5e-06, 'epoch': 0.98}\n",
      "{'loss': 0.0073, 'grad_norm': 0.29672884941101074, 'learning_rate': 2.0000000000000003e-06, 'epoch': 0.98}\n",
      "{'loss': 0.0022, 'grad_norm': 0.12063118070363998, 'learning_rate': 1.5e-06, 'epoch': 0.99}\n",
      "{'loss': 0.013, 'grad_norm': 0.35029494762420654, 'learning_rate': 1.0000000000000002e-06, 'epoch': 0.99}\n",
      "{'loss': 0.0043, 'grad_norm': 0.18529963493347168, 'learning_rate': 5.000000000000001e-07, 'epoch': 1.0}\n",
      "100%|█████████████████████████████████████████| 200/200 [02:02<00:00,  2.13it/s]\n",
      "  0%|                                                    | 0/50 [00:00<?, ?it/s]\u001b[A\n",
      "  4%|█▊                                          | 2/50 [00:00<00:06,  7.25it/s]\u001b[A\n",
      "  6%|██▋                                         | 3/50 [00:00<00:09,  5.13it/s]\u001b[A\n",
      "  8%|███▌                                        | 4/50 [00:00<00:10,  4.47it/s]\u001b[A\n",
      " 10%|████▍                                       | 5/50 [00:01<00:11,  4.06it/s]\u001b[A\n",
      " 12%|█████▎                                      | 6/50 [00:01<00:11,  3.83it/s]\u001b[A\n",
      " 14%|██████▏                                     | 7/50 [00:01<00:11,  3.74it/s]\u001b[A\n",
      " 16%|███████                                     | 8/50 [00:01<00:11,  3.66it/s]\u001b[A\n",
      " 18%|███████▉                                    | 9/50 [00:02<00:11,  3.62it/s]\u001b[A\n",
      " 20%|████████▌                                  | 10/50 [00:02<00:11,  3.59it/s]\u001b[A\n",
      " 22%|█████████▍                                 | 11/50 [00:02<00:10,  3.59it/s]\u001b[A\n",
      " 24%|██████████▎                                | 12/50 [00:03<00:10,  3.58it/s]\u001b[A\n",
      " 26%|███████████▏                               | 13/50 [00:03<00:10,  3.61it/s]\u001b[A\n",
      " 28%|████████████                               | 14/50 [00:03<00:09,  3.63it/s]\u001b[A\n",
      " 30%|████████████▉                              | 15/50 [00:03<00:09,  3.64it/s]\u001b[A\n",
      " 32%|█████████████▊                             | 16/50 [00:04<00:09,  3.63it/s]\u001b[A\n",
      " 34%|██████████████▌                            | 17/50 [00:04<00:09,  3.60it/s]\u001b[A\n",
      " 36%|███████████████▍                           | 18/50 [00:04<00:08,  3.57it/s]\u001b[A\n",
      " 38%|████████████████▎                          | 19/50 [00:05<00:08,  3.58it/s]\u001b[A\n",
      " 40%|█████████████████▏                         | 20/50 [00:05<00:08,  3.57it/s]\u001b[A\n",
      " 42%|██████████████████                         | 21/50 [00:05<00:08,  3.56it/s]\u001b[A\n",
      " 44%|██████████████████▉                        | 22/50 [00:05<00:07,  3.54it/s]\u001b[A\n",
      " 46%|███████████████████▊                       | 23/50 [00:06<00:07,  3.51it/s]\u001b[A\n",
      " 48%|████████████████████▋                      | 24/50 [00:06<00:07,  3.50it/s]\u001b[A\n",
      " 50%|█████████████████████▌                     | 25/50 [00:06<00:07,  3.51it/s]\u001b[A\n",
      " 52%|██████████████████████▎                    | 26/50 [00:07<00:06,  3.50it/s]\u001b[A\n",
      " 54%|███████████████████████▏                   | 27/50 [00:07<00:06,  3.51it/s]\u001b[A\n",
      " 56%|████████████████████████                   | 28/50 [00:07<00:06,  3.53it/s]\u001b[A\n",
      " 58%|████████████████████████▉                  | 29/50 [00:07<00:05,  3.52it/s]\u001b[A\n",
      " 60%|█████████████████████████▊                 | 30/50 [00:08<00:05,  3.49it/s]\u001b[A\n",
      " 62%|██████████████████████████▋                | 31/50 [00:08<00:05,  3.51it/s]\u001b[A\n",
      " 64%|███████████████████████████▌               | 32/50 [00:08<00:05,  3.53it/s]\u001b[A\n",
      " 66%|████████████████████████████▍              | 33/50 [00:09<00:04,  3.55it/s]\u001b[A\n",
      " 68%|█████████████████████████████▏             | 34/50 [00:09<00:04,  3.54it/s]\u001b[A\n",
      " 70%|██████████████████████████████             | 35/50 [00:09<00:04,  3.54it/s]\u001b[A\n",
      " 72%|██████████████████████████████▉            | 36/50 [00:09<00:03,  3.54it/s]\u001b[A\n",
      " 74%|███████████████████████████████▊           | 37/50 [00:10<00:03,  3.54it/s]\u001b[A\n",
      " 76%|████████████████████████████████▋          | 38/50 [00:10<00:03,  3.55it/s]\u001b[A\n",
      " 78%|█████████████████████████████████▌         | 39/50 [00:10<00:03,  3.56it/s]\u001b[A\n",
      " 80%|██████████████████████████████████▍        | 40/50 [00:10<00:02,  3.56it/s]\u001b[A\n",
      " 82%|███████████████████████████████████▎       | 41/50 [00:11<00:02,  3.57it/s]\u001b[A\n",
      " 84%|████████████████████████████████████       | 42/50 [00:11<00:02,  3.57it/s]\u001b[A\n",
      " 86%|████████████████████████████████████▉      | 43/50 [00:11<00:01,  3.57it/s]\u001b[A\n",
      " 88%|█████████████████████████████████████▊     | 44/50 [00:12<00:01,  3.56it/s]\u001b[A\n",
      " 90%|██████████████████████████████████████▋    | 45/50 [00:12<00:01,  3.52it/s]\u001b[A\n",
      " 92%|███████████████████████████████████████▌   | 46/50 [00:12<00:01,  3.52it/s]\u001b[A\n",
      " 94%|████████████████████████████████████████▍  | 47/50 [00:12<00:00,  3.52it/s]\u001b[A\n",
      " 96%|█████████████████████████████████████████▎ | 48/50 [00:13<00:00,  3.53it/s]\u001b[A\n",
      "                                                                                \n",
      "\u001b[A{'eval_loss': 0.01484476774930954, 'eval_runtime': 14.1016, 'eval_samples_per_second': 28.366, 'eval_steps_per_second': 3.546, 'epoch': 1.0}\n",
      "100%|█████████████████████████████████████████| 200/200 [02:16<00:00,  2.13it/s]\n",
      "100%|███████████████████████████████████████████| 50/50 [00:13<00:00,  3.55it/s]\u001b[A\n",
      "{'train_runtime': 139.1928, 'train_samples_per_second': 11.495, 'train_steps_per_second': 1.437, 'train_loss': 0.21926165691111238, 'epoch': 1.0}\n",
      "100%|█████████████████████████████████████████| 200/200 [02:19<00:00,  1.44it/s]\n",
      "TrainOutput(global_step=200, training_loss=0.21926165691111238, metrics={'train_runtime': 139.1928, 'train_samples_per_second': 11.495, 'train_steps_per_second': 1.437, 'total_flos': 5646339091599360.0, 'train_loss': 0.21926165691111238, 'epoch': 1.0})\n",
      "Peak VRAM during training: 17.510 GiB\n"
     ]
    }
   ],
   "source": [
    "!mdl-toolkit train \\\n",
    "     --lora-rank 32 \\\n",
    "     --eval-steps 100 \\\n",
    "     --train-dataset train-converted/ \\\n",
    "     --eval-dataset test-converted/ \\\n",
    "     --output output/ \\\n",
    "     --model-name mispeech/midashenglm-7b-bf16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e33BJPrBHa5N"
   },
   "source": [
    "## Inference\n",
    "\n",
    "After training, use the fine-tuned model for inference. By default, LoRA adapters are merged, so you can load it the same way as the base model by specifying the model path:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OiAIrMtWHa5N",
    "outputId": "58bd891b-20a9-4946-d082-2682e79250d0"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/h30/zhoujiahao5/notebook/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section'}\n",
      "Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:15<00:00,  3.78s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['category: pig, target: 2']\n"
     ]
    }
   ],
   "source": [
    "# Set TOKENIZERS_PARALLELISM=0 in the Notebook to prevent warnings from\n",
    "# messing up the output. Do not set this outside the Notebook, as it may\n",
    "# cause reduced performance.\n",
    "import os\n",
    "\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"0\"\n",
    "\n",
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer\n",
    "\n",
    "model_id = \"./output/final/\"\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)\n",
    "\n",
    "model.eval()\n",
    "\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"system\",\n",
    "        \"content\": [\n",
    "            {\n",
    "                \"type\": \"text\",\n",
    "                \"text\": \"Output the predicted category in the format of category: <category>, category_id: <category_id>.\",\n",
    "            },\n",
    "        ],\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"audio\", \"path\": \"ESC-50-master/audio/5-103415-A-2.wav\"},\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "with torch.no_grad():\n",
    "    model_inputs = processor.apply_chat_template(\n",
    "        messages,\n",
    "        tokenize=True,\n",
    "        add_generation_prompt=True,\n",
    "        add_special_tokens=True,\n",
    "        return_dict=True,\n",
    "    ).to(device=model.device, dtype=model.dtype)\n",
    "    generation = model.generate(**model_inputs)\n",
    "    output = tokenizer.batch_decode(generation, skip_special_tokens=True)\n",
    "\n",
    "print(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7OyfJdszHa5N"
   },
   "source": [
    "We can also use the MDL-Toolkit inference command to obtain predictions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "aEHBYrH3Ha5N",
    "outputId": "9b0b09ae-ae86-4cf6-a2af-e609e4fb2970"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section'}\n",
      "Loading checkpoint shards: 100%|██████████████████| 4/4 [00:04<00:00,  1.24s/it]\n",
      "Processing dataset: 100%|██████████████| 400/400 [00:10<00:00, 37.82 examples/s]\n",
      "Batching examples (num_proc=32): 100%|█| 400/400 [00:03<00:00, 119.73 examples/s\n",
      "Inference: 100%|████████████████████████████████| 32/32 [00:33<00:00,  1.06s/it]\n"
     ]
    }
   ],
   "source": [
    "!mdl-toolkit inference \\\n",
    "    test.csv \\\n",
    "    --system-prompt \"You are a helpful audio classifier.\" \\\n",
    "    --user-prompt \"Output the predicted category in the format of category: <category>, category_id: <category_id>.\" \\\n",
    "    --output finetuned-output.csv \\\n",
    "    --model-name ./output/final/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WrVClf7EHa5O"
   },
   "source": [
    "The predictions are written to the specified CSV file with the same schema as training. After fine-tuning, the model should produce outputs that follow the specified format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1zQKD5j8Ha5O",
    "outputId": "35241aa4-e074-4e7b-f055-7a01cca0cb19"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "audio,prediction\r\r\n",
      "ESC-50-master/audio/5-103415-A-2.wav,\"category: livestock, category_id: 1\"\r\r\n",
      "ESC-50-master/audio/5-103416-A-2.wav,\"category: music, category_id: 1\"\r\r\n",
      "ESC-50-master/audio/5-103418-A-2.wav,\"category: pig, category_id: 1\"\r\r\n",
      "ESC-50-master/audio/5-103420-A-2.wav,\"category: animal, category_id: 1\"\r\r\n",
      "ESC-50-master/audio/5-103421-A-2.wav,\"category: pig, category_id: 1\"\r\r\n",
      "ESC-50-master/audio/5-103422-A-2.wav,\"category: animal, category_id: 1\"\r\r\n",
      "ESC-50-master/audio/5-117118-A-42.wav,\"category: alarm, category_id: 1\"\r\r\n",
      "ESC-50-master/audio/5-117120-A-42.wav,\"category: alarm, category_id: 1\"\r\r\n",
      "ESC-50-master/audio/5-117122-A-42.wav,\"category: alarm, category_id: 1\"\r\r\n",
      "ESC-50-master/audio/5-117250-A-2.wav,\"category: animal, category_id: 1\"\n",
      "audio,prediction\n",
      "ESC-50-master/audio/5-103415-A-2.wav,\"category: pig, target: 2\"\n",
      "ESC-50-master/audio/5-103416-A-2.wav,\"category: door_wood_creaks, target: 33\"\n",
      "ESC-50-master/audio/5-103418-A-2.wav,\"category: pig, target: 2\"\n",
      "ESC-50-master/audio/5-103420-A-2.wav,\"category: rooster, target: 1\"\n",
      "ESC-50-master/audio/5-103421-A-2.wav,\"category: pig, target: 2\"\n",
      "ESC-50-master/audio/5-103422-A-2.wav,\"category: pig, target: 3\"\n",
      "ESC-50-master/audio/5-117118-A-42.wav,\"category: siren, target: 42\"\n",
      "ESC-50-master/audio/5-117120-A-42.wav,\"category: siren, target: 42\"\n",
      "ESC-50-master/audio/5-117122-A-42.wav,\"category: siren, target: 42\"\n",
      "ESC-50-master/audio/5-117250-A-2.wav,\"category: snoring, target: 28\"\n"
     ]
    }
   ],
   "source": [
    "! head -n 11 orig-output.csv\n",
    "! head -n 11 finetuned-output.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vU4pwn2jZWmB"
   },
   "source": [
    "## How to improve performance\n",
    "\n",
    "Congrats on completing your first fine-tune! Hyperparameters in this tutorial may not be optimal for all tasks. If you’re not satisfied with the results, try the following:\n",
    "\n",
    "1. Increase the LoRA rank, e.g., `--lora-rank 64`.\n",
    "2. Tune the learning rate, e.g., `--lr 5e-5`. The best LR depends on many factors and may require multiple trials or a systematic search.\n",
    "3. Adjust trainable targets, e.g., `--train-target encoder --train-target projector --train-target decoder --train-target embed_tokens --train-target lm_head` to train all available targets. In some cases, adding targets—especially `embed_tokens` and `lm_head`—can improve results.\n",
    "4. Use higher numerical precision. If you used quantization, try running without it. If you didn’t, you can set `--bf16 false` to load fp32 weights.\n",
    "5. Increase both the quantity and quality of training data. This often helps, although naively repeating the same data (e.g., setting `--num-epochs` > 1) may have limited effect or even hurt performance."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}